package com.kara.woodAgent.agent.wood;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSON;
import com.kara.woodAgent.agent.graph.BaseGraph;
import com.kara.woodAgent.agent.graph.Next;
import com.kara.woodAgent.agent.graph.impl.GraphNodeImpl;
import com.kara.woodAgent.agent.model.ModelProvider;
import com.kara.woodAgent.agent.tool.provider.ToolProvider;
import com.kara.woodAgent.agent.wood.context.WoodContext;
import com.kara.woodAgent.agent.wood.context.WoodReactContext;
import com.kara.woodAgent.agent.wood.model.ExpectedMessage;
import com.kara.woodAgent.agent.wood.model.Plan;
import com.kara.woodAgent.agent.wood.model.SpecialAgentResult;
import com.kara.woodAgent.agent.wood.node.ActionToolNode;
import com.kara.woodAgent.agent.wood.node.ReviewNode;
import com.kara.woodAgent.agent.wood.node.SummaryNode;
import com.kara.woodAgent.agent.wood.node.ThinkNode;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;

import java.util.ArrayList;
import java.util.List;

public class SpecialAgentTemplate {


    public static GraphNodeImpl<WoodContext, WoodReactContext> createGraphNode(ModelProvider modelProvider, ToolProvider.ToolServiceHolder toolServiceHolder) {
        //正常情况下 lastMessage是 退出原因

        BaseGraph<WoodReactContext> reactGraph = new BaseGraph<WoodReactContext>()
                .addStartNode("model", new ThinkNode(modelProvider, toolServiceHolder))
                .addNode("review",new ReviewNode(modelProvider, toolServiceHolder))
                .addNode("tool", new ActionToolNode(toolServiceHolder))
                .addNode("summary", new SummaryNode(modelProvider));

        return new GraphNodeImpl<WoodContext, WoodReactContext>(reactGraph)
                .gcSupplier(woodContext -> {

                    Plan currentPlan = woodContext.getCurrentPlan();

                    String subQuestion = currentPlan.getContent();

                    woodContext.setChatMessages(null);
                    String systemPrompt;

                    List<ChatMessage>  chatMessages = new ArrayList<>();

                    systemPrompt = """
                                #### 身份
                                你是wood智能体系统 是由 wood团队开发的通用智能体系统.
                                #### 职责
                                    结合执行的历史信息 和 分配给你的任务,完成分配给你的任务
                                来完成任务.
                                #### 执行的历史信息包括 ####
                                {{history}}
                                #### 你的任务是 ####
                                {{question}}
                                """;


                    SystemMessage systemMessage = SystemMessage.from(systemPrompt.replace("{{question}}", subQuestion)
                            .replace("{{history}}", (woodContext.getHistory() != null && !woodContext.getHistory().isEmpty()) ? JSON.toJSONString(woodContext.getHistory()) : "无历史消息"));
                    chatMessages.add(systemMessage);

                    WoodReactContext reactGraphContext = new WoodReactContext();
                    ExpectedMessage expectedMessage = new ExpectedMessage(chatMessages);
                    reactGraphContext.setExpectedMessage(expectedMessage);
                    reactGraphContext.setPlan(subQuestion);
                    reactGraphContext.setChatMessages(chatMessages);

                    return reactGraphContext;
                })
                .next((woodContext, reactGraphContext) -> {

                    //正常情况下 lastMessage是 退出原因
                    SpecialAgentResult result = reactGraphContext.getSummary();

                    Plan currentPlan = woodContext.getCurrentPlan();
                    currentPlan.setStatus(result.status());
                    if (Plan.SUCCESS.equals(currentPlan.getStatus())) {
                        currentPlan.setResultMsg(result.message());
                    }
                    List<Plan> plans = woodContext.getPlans();
                    List<String> list = plans.stream().map(Plan::getResultMsg).filter(StrUtil::isNotBlank).toList();
                    woodContext.setHistory(list);
                    List<ChatMessage> chatMessages = reactGraphContext.getChatMessages();

                    List<ChatMessage> messages = woodContext.getChatMessages();
                    if(messages == null){
                        messages=new ArrayList<>();
                    }

                    messages.addAll(chatMessages);
                    woodContext.setChatMessages(messages);

                    return Next.NextNode("dispatch").step(JSON.toJSONString(currentPlan));
                });
    }


}
